
# Function to do MCMC sampling
MyCAR <- function(Y,adjs,X,N,firstday=1,
                  nugget=TRUE,spatial=TRUE,temporal=TRUE,
                  iters=200,burn=100,thin=1,update=1){
   library(spam)
   # Set up
   ns         <- nrow(Y) 
   ms         <- unlist(lapply(adjs,length))
   adj1       <- rep(1:ns,times=ms)
   adj2       <- unlist(adjs)

   # Set up time adjacency
   nt         <- ncol(Y)
   mt         <- c(1,rep(2,nt-2),1)
   adjt       <- list()
   adjt[[1]]  <- 2
   for(t in 2:(nt-1)){adjt[[t]]<-c(t-1,t+1)}
   adjt[[nt]] <- nt-1

   if(!is.matrix(N)){N <- matrix(N,ns,nt)}

   # Computation for the determinant
   ADJt  <- as.spam(as.matrix(dist(1:nt))==1)
   ADJs <- as.spam(matrix(0,ns,ns))
   for(s in 1:ns){ADJs[s,adjs[[s]]]<-1}
   Cs <- sweep(ADJs,1,sqrt(ms),"/")
   Cs <- sweep(Cs,2,sqrt(ms),"/")
   Ct <- sweep(ADJt,1,sqrt(mt),"/")
   Ct <- sweep(Ct,2,sqrt(mt),"/")
   ds <- eigen(Cs)$val
   dt <- eigen(Ct)$val
   rm(Cs,Ct)

  # Initial values

  # lambda(t) = N*[exp(e)+exp(v)]
  # v(t) ~iid N(muv, tauv)
  # e(t) ~indep N(theta(t), taue)
  # theta = mn_theta(alpha) + STCAR

   p     <- length(X)
   beta  <- rep(0,p)
   taus  <- 1
   taue  <- 10
   tauv  <- 10
   muv   <- -2
   rhos  <- ifelse(spatial,0.95,0)
   rhot  <- ifelse(temporal,0.50,0)
   e     <- log((Y+10)/(N+10))
   theta <- 0.95*e + 0.05*mean(e)
   v     <- ifelse(nugget,-2,-Inf)*matrix(1,ns,nt)
   
   # Indictors of which observations should be used
   IN <- matrix(0,ns,nt)
   IN[,firstday:nt]<-1
   Y  <- IN*Y
   N  <- IN*N

   mu <- 0
   for(j in 1:p){mu <- mu + X[[j]]*beta[j]}

   MH <- att <- acc <- c(2,2,0.05,0.05,0.05,0.05)
   M1 <- thetahat1 <- thetahat2 <- 0
   
   keepers <- matrix(0,iters,p+6)
   dev     <- rep(0,iters)
   colnames(keepers)<-c(paste0("beta",1:p),
                        "taue","taus","rhos","rhot",
                        "muv","tauv")
              
  tick <- proc.time()[3]
  for(iter in 1:iters){
    for(ttt in 1:thin){

    # Update nugget effect, e

      cane <- e+MH[1]*rnorm(ns*nt)/sqrt(taue)
      R    <- dpois(Y,N*(exp(cane)+exp(v)),log=TRUE)-
              dpois(Y,N*(exp(e)+exp(v)),log=TRUE)+
              dnorm(cane,theta,1/sqrt(taue),log=TRUE)-
              dnorm(e,theta,1/sqrt(taue),log=TRUE)
      R[is.na(R)] <- -Inf
      accept      <- log(runif(ns*nt))<R
      e[accept]   <- cane[accept] 
      acc[1]      <- acc[1] + sum(accept)
      att[1]      <- att[1] + ns*nt

    if(nugget){
       canv <- v+MH[2]*rnorm(ns*nt)/sqrt(tauv)
       R    <- dpois(Y,N*(exp(e)+exp(canv)),log=TRUE)-
               dpois(Y,N*(exp(e)+exp(v)),log=TRUE)+
               dnorm(canv,muv,1/sqrt(tauv),log=TRUE)-
               dnorm(v,muv,1/sqrt(tauv),log=TRUE)
       R[is.na(R)] <- -Inf
       accept      <- log(runif(ns*nt))<R
       v[accept]   <- canv[accept] 
       acc[2]      <- acc[2] + sum(accept)
       att[2]      <- att[2] + ns*nt
    } 

    taue <- rgamma(1,ns*nt/2+0.1,sum((e-theta)^2)/2+0.1)
    if(nugget){
      tauv <- rgamma(1,ns*nt/2+0.1,sum((v-muv)^2)/2+0.1)
      VVV  <- tauv*ns*nt+0.01
      MMM  <- tauv*sum(v)
      muv  <- rnorm(1,MMM/VVV,1/sqrt(VVV))
    }  


    # Update spatiotemporal process theta
    R   <- theta-mu
    Qt  <- taus*(diag(mt)-rhot*ADJt)
    for(j in unique(ms)){
      VVV <- solve(j*Qt + taue*diag(nt))
      PPP <- t(chol(VVV))
      for(k in which(ms==j)){
        MMM <- R[adjs[[k]],]
        if(j>1){MMM <- colSums(MMM)}
        MMM       <- Qt%*%(j*mu[k,]+rhos*MMM) + taue*e[k,]
        theta[k,] <- VVV%*%MMM + PPP%*%rnorm(nt)
        R[k,]     <- theta[k,]-mu[k,]
      }
    }

    SS   <- qf(R,rhos,rhot,ns,nt,ms,mt,adj1,adj2)
    taus <- rgamma(1,ns*nt/2+0.1,SS/2+0.1)

    # Do updates on the z scale
    if(spatial){
      att[5]  <- att[5] + 1
      z       <- qnorm(pbeta(rhos,1,1))
      canz    <- rnorm(1,z,MH[5])
      canrhos <- qbeta(pnorm(canz),1,1)
      canSS   <- qf(R,canrhos,rhot,ns,nt,ms,mt,adj1,adj2)
      MHR     <- 0.5*logdet(canrhos,rhot,ds,dt) - 0.5*taus*canSS -
                 0.5*logdet(rhos,rhot,ds,dt) + 0.5*taus*SS +
                 dnorm(canz,log=TRUE)-dnorm(z,log=TRUE)
      if(canrhos<0.995){if(log(runif(1))<MHR){
           rhos<-canrhos;SS<-canSS;acc[5] <- acc[5] + 1
      }}   
    }
    if(temporal){ 
      att[6]  <- att[6] + 1
      z       <- qnorm(pbeta(rhot,1,1))
      canz    <- rnorm(1,z,MH[6])
      canrhot <- qbeta(pnorm(canz),1,1)
      canSS   <- qf(R,rhos,canrhot,ns,nt,ms,mt,adj1,adj2)
      MHR     <- 0.5*logdet(rhos,canrhot,ds,dt) - 0.5*taus*canSS -
                 0.5*logdet(rhos,rhot,ds,dt) + 0.5*taus*SS +
                 dnorm(canz,log=TRUE)-dnorm(z,log=TRUE)
      if(log(runif(1))<MHR){
           rhot<-canrhot;SS<-canSS;acc[6] <- acc[6] + 1
      }   
    }


   # Update beta
    Qt   <- taus*(diag(mt)-rhot*ADJt)
    L    <- chol(Qt)
    for(j in 1:p){
       A       <- X[[j]]
       mu      <- mu - A*beta[j]
       La      <- L%*%t(A)
       Lt      <- L%*%t(theta-mu)
       VVV     <- 0.01 + sum(ms*colSums(La*La))-rhos*sum(La[,adj1]*La[,adj2])
       MMM     <- 0    + sum(ms*colSums(La*Lt))-rhos*sum(La[,adj1]*Lt[,adj2])
       beta[j] <- rnorm(1,MMM/VVV,1/sqrt(VVV))
       mu      <- mu + A*beta[j]
    }



    } # end thin


    # Tuning

    if(iter<burn){for(j in 1:length(MH)){if(att[j]>25){
       if(acc[j]/att[j]<0.3){MH[j]<-MH[j]*0.8}
       if(acc[j]/att[j]>0.5){MH[j]<-MH[j]*1.2}
       acc[j]<-att[j]<-0
    }}}

    # Keep track of output
    keepers[iter,] <- c(beta,taue,taus,rhos,rhot,muv,tauv)
    MM             <- exp(e)+exp(v)
    dev[iter]      <- -2*sum(dpois(Y,N*MM,log=TRUE))
    if(iter>burn){
      thetahat1 <- thetahat1+theta/(iters-burn)
      thetahat2 <- thetahat2+theta*theta/(iters-burn)
      M1        <- M1 + MM/(iters-burn)
    }
 
    if(iter%%update==0){
      par(mfrow=c(3,2))
      for(j in 2:7){
          plot(keepers[1:iter,j],main=colnames(keepers)[j],type="l")
          abline(0.2,0,lty=2);abline(0.5,0,lty=2)
      }
    }

  } # end iters

  dbar <- mean(dev[burn:iters])
  dhat <- -2*sum(dpois(Y,N*M1,log=TRUE))
  pD   <- dbar-dhat
  DIC  <- dbar + pD

  tock <- proc.time()[3]
  out <- list(keepers=keepers,
              theta_hat = thetahat1,
              theta_var = thetahat2-thetahat1^2,
              min = (tock-tick)/60,
              DIC = list(DIC=DIC,dbar=dbar,pD=pD),
              dev=dev)
return(out)}

# Det of the spatiotemporal CAR covariance
logdet <- function(rhos,rhot,ds,dt){
   length(ds)*sum(log(1-rhot*dt)) +
   length(dt)*sum(log(1-rhos*ds))
}

# Compute CAR quadratic form
# t(R)%*%[SigmaS(rhos)\kroneckerSigmaT(rhot)]%*%R
qf <- function(R,rhos,rhot,ns,nt,ms,mt,adj1,adj2){
  SS <- ms%*%(R^2)%*%mt-rhos*sum((R[adj1,]*R[adj2,])%*%mt)
  SS <- SS-rhot*sum(ms%*%(R[,-1]*R[,-nt]))
  SS <- SS+rhot*rhos*sum(R[adj1,-1]*R[adj2,-nt])
return(as.vector(SS))}


# Compute the weighted function for the reporting lag
make.w <- function(nt,lag,bw){
  # w[j,t] = weight of lambda[j] on Z[t]
  w <- matrix(0,nt,nt)
  for(j in 1:nt){
    ww    <- (1:nt-j-lag)/bw
    w[j,] <- ifelse(1:nt<j,0,exp(-ww^2))
  } 
return(w)}

